from fastapi import FastAPI, HTTPException, Query
from pydantic import BaseModel
from typing import Optional, Dict, List
import uvicorn
from TCMLLMHandler import TCMLLM

# 创建FastAPI应用实例
app = FastAPI(
    title="中医药知识库API",
    description="基于中医药数据库的知识查询API，支持药材信息、功效、主治等查询",
    version="1.0.0",
)

# 初始化TCMLLM实例
tcm_llm = TCMLLM("tcm.db")


# 定义响应模型
class HerbResponse(BaseModel):
    herb_name: str
    alias: Optional[str] = None
    property: Optional[str] = None
    meridian: Optional[str] = None
    efficacy: Optional[str] = None
    indication: Optional[str] = None
    usage: Optional[str] = None
    contraindication: Optional[str] = None


# 定义错误响应模型
class ErrorResponse(BaseModel):
    error: str
    detail: Optional[str] = None


# 根路由 - API欢迎信息
@app.get("/", tags=["首页"])
async def root():
    return {"message": "欢迎使用中医药知识库API！", "version": "1.0.0"}


# 药材查询路由
@app.get(
    "/herb/{herb_name}",
    response_model=HerbResponse,
    tags=["药材查询"],
    responses={404: {"model": ErrorResponse, "description": "未找到药材"}},
)
async def get_herb(herb_name: str):
    """根据药材名称查询完整信息"""
    herb_data = tcm_llm.db.get_herb_by_name(herb_name)
    if not herb_data:
        # 尝试通过别名查询
        if herb_name in tcm_llm.alias_map:
            herb_data = tcm_llm.db.get_herb_by_name(tcm_llm.alias_map[herb_name])

    if not herb_data:
        raise HTTPException(status_code=404, detail=f"未找到药材【{herb_name}】的信息")

    return herb_data


# 药材属性查询路由
@app.get(
    "/herb/{herb_name}/{property_type}",
    tags=["药材属性查询"],
    responses={404: {"model": ErrorResponse, "description": "未找到药材或属性"}},
)
async def get_herb_property(herb_name: str, property_type: str):
    """查询药材的特定属性（功效、主治、性味、归经等）"""
    valid_properties = {
        "efficacy",
        "indication",
        "property",
        "meridian",
        "usage",
        "contraindication",
    }

    if property_type not in valid_properties:
        raise HTTPException(
            status_code=400,
            detail=f"无效的属性类型！有效值：{', '.join(valid_properties)}",
        )

    # 处理别名
    if herb_name in tcm_llm.alias_map:
        herb_name = tcm_llm.alias_map[herb_name]

    herb_data = tcm_llm.db.get_herb_by_name(herb_name)
    if not herb_data:
        raise HTTPException(status_code=404, detail=f"未找到药材【{herb_name}】的信息")

    return {
        "herb_name": herb_name,
        property_type: herb_data.get(property_type, "无相关记载"),
    }


# 按功效搜索药材
@app.get("/search/", tags=["药材搜索"])
async def search_herbs(
    keyword: str = Query(..., description="搜索关键词"),
    property_type: str = Query(
        "efficacy",
        description="属性类型：efficacy(功效), indication(主治), property(性味), meridian(归经)",
    ),
):
    """按药性属性（功效、主治、性味、归经）搜索药材"""
    valid_properties = {"efficacy", "indication", "property", "meridian"}

    if property_type not in valid_properties:
        raise HTTPException(
            status_code=400,
            detail=f"无效的属性类型！有效值：{', '.join(valid_properties)}",
        )

    results = tcm_llm.db.search_herb_by_property(property_type, keyword)

    if not results:
        return {"message": f"未找到与【{keyword}】相关的药材", "herbs": []}

    # 简化返回结果，只包含药材名称和匹配的属性
    simplified_results = [
        {"herb_name": herb["herb_name"], property_type: herb[property_type]}
        for herb in results
    ]

    return {"message": f"找到{len(results)}个相关药材", "herbs": simplified_results}


# 自然语言查询接口
@app.get("/query/", tags=["自然语言查询"])
async def natural_language_query(
    q: str = Query(..., description="自然语言问题，如：黄芪的功效是什么？")
):
    """处理自然语言形式的中药查询"""
    response = tcm_llm.process_query(q)
    return {"query": q, "response": response}


# 启动服务器（当直接运行此文件时）
if __name__ == "__main__":
    uvicorn.run("api:app", host="0.0.0.0", port=8000, reload=True)
